import os
import tempfile
from seal import *
from app.pkg.snark_verifier import SNARKVerifier
from app.pkg.utils import *
from collections import namedtuple

QRP = namedtuple('QRP', ['V', 'W', 'Y', 't'])


class FileDecryptor:
    def __init__(self, context, key_manager, oss_handler):
        self.context = context
        self.key_manager = key_manager
        self.oss_handler = oss_handler
        self._verify_cache = {}

    # def decrypt_file(self, oss_dir: str, output_path: str):
    def decrypt_file(self, oss_dir: str):
        try:
            # 创建临时目录用于存储下载的文件
            temp_dir = tempfile.mkdtemp()
            encrypted_folder = os.path.join(temp_dir, 'encrypted_files')
            evidence_folder = os.path.join(temp_dir, 'evidence_files')
            os.makedirs(encrypted_folder, exist_ok=True)
            os.makedirs(evidence_folder, exist_ok=True)

            # 下载并加载密钥
            private_key_path = os.path.join(temp_dir, 'private_key.key')
            public_key_path = os.path.join(temp_dir, 'public_key.key')
            self.key_manager.download_private_key(oss_dir + 'keys/private_key.key', private_key_path)
            self.key_manager.download_public_key(oss_dir + 'keys/public_key.key', public_key_path)

            secret_key = SecretKey()
            public_key = PublicKey()
            secret_key.load(self.context, private_key_path)
            public_key.load(self.context, public_key_path)

            decryptor = Decryptor(self.context, secret_key)
            decoder = BatchEncoder(self.context)

            # 下载和加载 CRS 和 VRS
            crs_file = os.path.join(evidence_folder, 'evidence.crs')
            vrs_file = os.path.join(evidence_folder, 'evidence.vrs')
            self.oss_handler.download_file(oss_dir + 'evidence.crs', crs_file)
            self.oss_handler.download_file(oss_dir + 'evidence.vrs', vrs_file)

            with open(crs_file, 'r') as f:
                crs_data = json.load(f)
                h = np.poly1d(crs_data['h'])
                qrp = QRP(
                    [np.poly1d(v) for v in crs_data['qrp']['V']],
                    [np.poly1d(w) for w in crs_data['qrp']['W']],
                    [np.poly1d(y) for y in crs_data['qrp']['Y']],
                    np.poly1d(crs_data['qrp']['t'])
                )
                verifier = SNARKVerifier(self.context)
                verifier.h = h
                verifier.crs = (h, qrp, crs_data['crs_data'])

            with open(vrs_file, 'r') as f:
                verifier.vrs = (h, qrp, json.load(f))

            # 下载 chunk_info.json
            chunk_info_file = os.path.join(encrypted_folder, 'chunk_info.json')
            self.oss_handler.download_file(oss_dir + 'chunk_info.json', chunk_info_file)

            decrypted_content = []
            with open(chunk_info_file, 'r') as f:
                chunk_info = json.load(f)
                for chunk in chunk_info:
                    chunk_file_name = chunk["file"]
                    length = chunk["length"]
                    local_chunk_path = os.path.join(encrypted_folder, os.path.basename(chunk_file_name))

                    if not self.oss_handler.download_file(oss_dir + chunk_file_name, local_chunk_path):
                        raise FileNotFoundError(f"Missing chunk file: {chunk_file_name}")

                    compressed = Ciphertext()
                    compressed.load(self.context, local_chunk_path)

                    proof = verifier.generate_proof(compressed, None)
                    if not verifier.verify(proof, compressed):
                        raise ValueError(f"Integrity check failed for {chunk_file_name}")

                    decrypted_plaintext = Plaintext()
                    decryptor.decrypt(compressed, decrypted_plaintext)
                    decoded_ints = decoder.decode(decrypted_plaintext)
                    decrypted_content.extend(decoded_ints[:length])

            # 下载并拼接 header 和内容
            header_file = os.path.join(encrypted_folder, 'header.bin')
            # if self.oss_handler.download_file(oss_dir + 'header.bin', header_file):
            #     with open(header_file, 'rb') as h:
            #         header = h.read()
            #     with open('dd.jpg', 'wb') as out_f:
            #         out_f.write(header)
            #         out_f.write(bytes(decrypted_content))
            # else:
            #     with open('dd.jpg', 'wb') as out_f:
            #         out_f.write(bytes(decrypted_content))
            # print("Decryption completed successfully.")
            full_data = b''
            if self.oss_handler.download_file(oss_dir + 'header.bin', header_file):
                with open(header_file, 'rb') as h:
                    header = h.read()
                full_data = header + bytes(decrypted_content)
            else:
                full_data = bytes(decrypted_content)

            print("✅ 解密完成，返回字节流")
            return full_data

        except Exception as e:
            print(f"[解密失败] Decryption error: {e}")
            raise
